package dao

import (
	"errors"
	"fmt"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/mysql"
	"github.com/spf13/viper"
	"math/rand"
	"sync"
	"time"
)

var (
	MasterDB    *gorm.DB
	MsterDBPool *DBPool
)

type DBPool struct {
	pool []*gorm.DB
	mu   *sync.RWMutex
}

// 初始化数据连接池
func (dbp *DBPool) init(num int) error {
	var (
		once       = new(sync.Once)
		atLeastOne = make(chan struct{})
		timeout    = time.NewTicker(30 * time.Second)
	)
	for i := 0; i < num; i++ {
		go func() {
			var (
				err     error
				temConn *gorm.DB
			)
			dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=%s&parseTime=True&loc=Local",
				viper.GetString("storage.user"),
				viper.GetString("storage.password"),
				viper.GetString("storage.host"),
				viper.GetString("storage.port"),
				viper.GetString("storage.dbname"),
				viper.GetString("storage.charset"))
			temConn, err = gorm.Open("mysql", dsn)
			if err != nil {
				fmt.Println("打开一个连接数据库失败: ", err)
				return
			}
			dbp.mu.Lock()
			defer dbp.mu.Unlock()
			dbp.pool = append(dbp.pool, temConn)
			once.Do(func() {
				atLeastOne <- struct{}{}
			})
		}()
	}
	select {
	case <-atLeastOne:
		return nil
	case <-timeout.C:
		return errors.New("超时并一个连接也没有成功")
	}
}

// 获取一个数据库连接
func (dbp *DBPool) GetDB() *gorm.DB {
	var (
		dbNum int
		idx   int
	)
	dbp.mu.RLock()
	defer dbp.mu.RUnlock()
	dbNum = len(dbp.pool)
	idx = rand.Intn(dbNum)
	return dbp.pool[idx]
}

// 创建一个数据库连接池
func CreateDBPool(num int) (*DBPool, error) {
	var (
		err error
	)
	newDBPool := DBPool{
		pool: []*gorm.DB{},
		mu:   new(sync.RWMutex),
	}
	err = newDBPool.init(num)
	if err != nil {
		return nil, err
	}
	return &newDBPool, nil
}

func Init() {
	var (
		err error
	)
	viper.SetDefault("storage.user", "root")
	viper.SetDefault("storage.password", "")
	viper.SetDefault("storage.host", "localhost")
	viper.SetDefault("storage.port", "3306")
	viper.SetDefault("storage.dbname", "crawler")
	viper.SetDefault("storage.charset", "utf8mb4")
	viper.SetDefault("storage.poolSize", "1")
	// 初始化数据库连接池
	MsterDBPool, err = CreateDBPool(viper.GetInt("storage.poolSize"))
	if err != nil {
		panic(err)
	}
	// 第一个连接为默认连接
	MasterDB = MsterDBPool.GetDB()
}
